[FIR] Fix exhaustiveness check in case of sealed subclass

The problem appear because check was performed only on direct children
of sealed class. To fix this, heritage tree is built and used in check

#KT-38989 Fixed

compiler/fir/analysis-tests/testData/resolve/exhaustiveness_sealedSubClass.kt

@@ -0,0 +1,54 @@
+sealed class A
+
+sealed class B : A()
+class C : A()
+
+sealed class D : B()
+sealed class E : B()
+
+fun test_1(e: A) {
+    val a = when (e) {
+        is C -> 1
+        is D -> 2
+        is E -> 3
+    }.plus(0)
+
+    val b = when (e) {
+        is B -> 1
+        is C -> 2
+    }.plus(0)
+
+    val c = when (e) {
+        is B -> 1
+        is C -> 2
+        is E -> 3
+        is D -> 4
+    }.plus(0)
+
+    val d = when (e) {
+        is E -> 1
+        is A -> 2
+    }.plus(0)
+}
+
+fun test_2(e: A) {
+    val a = when (e) {
+        is D -> 1
+        is E -> 2
+    }.<!UNRESOLVED_REFERENCE!>plus<!>(0)
+
+    val b = when (e) {
+        is B -> 1
+        is D -> 2
+        is E -> 3
+    }.<!UNRESOLVED_REFERENCE!>plus<!>(0)
+
+    val c = when (e) {
+        is B -> 1
+        is D -> 2
+    }.<!UNRESOLVED_REFERENCE!>plus<!>(0)
+
+    val d = when (e) {
+        is C -> 1
+    }.<!UNRESOLVED_REFERENCE!>plus<!>(0)
+}

compiler/fir/analysis-tests/testData/resolve/exhaustiveness_sealedSubClass.txt

@@ -0,0 +1,116 @@
+FILE: exhaustiveness_sealedSubClass.kt
+    public sealed class A : R|kotlin/Any| {
+        private constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public sealed class B : R|A| {
+        private constructor(): R|B| {
+            super<R|A|>()
+        }
+
+    }
+    public final class C : R|A| {
+        public constructor(): R|C| {
+            super<R|A|>()
+        }
+
+    }
+    public sealed class D : R|B| {
+        private constructor(): R|D| {
+            super<R|B|>()
+        }
+
+    }
+    public sealed class E : R|B| {
+        private constructor(): R|E| {
+            super<R|B|>()
+        }
+
+    }
+    public final fun test_1(e: R|A|): R|kotlin/Unit| {
+        lval a: R|kotlin/Int| = when (R|<local>/e|) {
+            ($subj$ is R|C|) ->  {
+                Int(1)
+            }
+            ($subj$ is R|D|) ->  {
+                Int(2)
+            }
+            ($subj$ is R|E|) ->  {
+                Int(3)
+            }
+        }
+        .R|kotlin/Int.plus|(Int(0))
+        lval b: R|kotlin/Int| = when (R|<local>/e|) {
+            ($subj$ is R|B|) ->  {
+                Int(1)
+            }
+            ($subj$ is R|C|) ->  {
+                Int(2)
+            }
+        }
+        .R|kotlin/Int.plus|(Int(0))
+        lval c: R|kotlin/Int| = when (R|<local>/e|) {
+            ($subj$ is R|B|) ->  {
+                Int(1)
+            }
+            ($subj$ is R|C|) ->  {
+                Int(2)
+            }
+            ($subj$ is R|E|) ->  {
+                Int(3)
+            }
+            ($subj$ is R|D|) ->  {
+                Int(4)
+            }
+        }
+        .R|kotlin/Int.plus|(Int(0))
+        lval d: R|kotlin/Int| = when (R|<local>/e|) {
+            ($subj$ is R|E|) ->  {
+                Int(1)
+            }
+            ($subj$ is R|A|) ->  {
+                Int(2)
+            }
+        }
+        .R|kotlin/Int.plus|(Int(0))
+    }
+    public final fun test_2(e: R|A|): R|kotlin/Unit| {
+        lval a: <ERROR TYPE REF: Unresolved name: plus> = when (R|<local>/e|) {
+            ($subj$ is R|D|) ->  {
+                Int(1)
+            }
+            ($subj$ is R|E|) ->  {
+                Int(2)
+            }
+        }
+        .<Unresolved name: plus>#(Int(0))
+        lval b: <ERROR TYPE REF: Unresolved name: plus> = when (R|<local>/e|) {
+            ($subj$ is R|B|) ->  {
+                Int(1)
+            }
+            ($subj$ is R|D|) ->  {
+                Int(2)
+            }
+            ($subj$ is R|E|) ->  {
+                Int(3)
+            }
+        }
+        .<Unresolved name: plus>#(Int(0))
+        lval c: <ERROR TYPE REF: Unresolved name: plus> = when (R|<local>/e|) {
+            ($subj$ is R|B|) ->  {
+                Int(1)
+            }
+            ($subj$ is R|D|) ->  {
+                Int(2)
+            }
+        }
+        .<Unresolved name: plus>#(Int(0))
+        lval d: <ERROR TYPE REF: Unresolved name: plus> = when (R|<local>/e|) {
+            ($subj$ is R|C|) ->  {
+                Int(1)
+            }
+        }
+        .<Unresolved name: plus>#(Int(0))
+    }

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -143,6 +143,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
         runTest("compiler/fir/analysis-tests/testData/resolve/exhaustiveness_sealedObject.kt");
     }
 
+    @TestMetadata("exhaustiveness_sealedSubClass.kt")
+    public void testExhaustiveness_sealedSubClass() throws Exception {
+        runTest("compiler/fir/analysis-tests/testData/resolve/exhaustiveness_sealedSubClass.kt");
+    }
+
     @TestMetadata("extension.kt")
     public void testExtension() throws Exception {
         runTest("compiler/fir/analysis-tests/testData/resolve/extension.kt");

compiler/fir/analysis-tests/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -143,6 +143,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
         runTest("compiler/fir/analysis-tests/testData/resolve/exhaustiveness_sealedObject.kt");
     }
 
+    @TestMetadata("exhaustiveness_sealedSubClass.kt")
+    public void testExhaustiveness_sealedSubClass() throws Exception {
+        runTest("compiler/fir/analysis-tests/testData/resolve/exhaustiveness_sealedSubClass.kt");
+    }
+
     @TestMetadata("extension.kt")
     public void testExtension() throws Exception {
         runTest("compiler/fir/analysis-tests/testData/resolve/extension.kt");

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirWhenExhaustivenessTransformer.kt

@@ -8,15 +8,11 @@ package org.jetbrains.kotlin.fir.resolve.transformers
 import org.jetbrains.kotlin.descriptors.ClassKind
 import org.jetbrains.kotlin.descriptors.Modality
 import org.jetbrains.kotlin.fir.FirElement
-import org.jetbrains.kotlin.fir.declarations.FirEnumEntry
+import org.jetbrains.kotlin.fir.declarations.*
-import org.jetbrains.kotlin.fir.declarations.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.collectEnumEntries
-import org.jetbrains.kotlin.fir.declarations.modality
 import org.jetbrains.kotlin.fir.expressions.*
 import org.jetbrains.kotlin.fir.expressions.impl.FirElseIfTrueCondition
 import org.jetbrains.kotlin.fir.references.FirResolvedNamedReference
 import org.jetbrains.kotlin.fir.resolve.BodyResolveComponents
-import org.jetbrains.kotlin.fir.resolve.providers.FirSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.firSymbolProvider
 import org.jetbrains.kotlin.fir.resolve.toSymbol
 import org.jetbrains.kotlin.fir.symbols.ConeClassLikeLookupTag
@@ -124,24 +120,49 @@ class FirWhenExhaustivenessTransformer(private val bodyResolveComponents: BodyRe
         sealedClass: FirRegularClass,
         nullable: Boolean
     ): Boolean {
-        val inheritors = sealedClass.sealedInheritors
+        if (sealedClass.sealedInheritors.isNullOrEmpty()) return true
-        if (inheritors.isNullOrEmpty()) return true
+
-        val data = SealedExhaustivenessData(
+        val data = SealedExhaustivenessData(sealedClass, !nullable)
-            sealedClass.session.firSymbolProvider,
-            inheritors.toMutableSet(),
-            !nullable
-        )
         for (branch in whenExpression.branches) {
             branch.condition.accept(SealedExhaustivenessVisitor, data)
         }
-        return data.containsNull && data.remainingInheritors.isEmpty()
+        return data.isExhaustive()
     }
 
-    private class SealedExhaustivenessData(
+    private inner class SealedExhaustivenessData(regularClass: FirRegularClass, var containsNull: Boolean) {
-        val symbolProvider: FirSymbolProvider,
+        val symbolProvider = bodyResolveComponents.symbolProvider
-        val remainingInheritors: MutableSet<ClassId>,
+        private val rootNode = SealedClassInheritors(regularClass.classId, regularClass.sealedInheritors.mapToSealedInheritors())
-        var containsNull: Boolean
+
-    )
+        private fun List<ClassId>?.mapToSealedInheritors(): MutableSet<SealedClassInheritors>? {
+            if (this.isNullOrEmpty()) return null
+
+            return this.mapNotNull {
+                val inheritor = symbolProvider.getClassLikeSymbolByFqName(it)?.fir as? FirRegularClass ?: return@mapNotNull null
+                SealedClassInheritors(inheritor.classId, inheritor.sealedInheritors.mapToSealedInheritors())
+            }.takeIf { it.isNotEmpty() }?.toMutableSet()
+        }
+
+        fun removeInheritor(classId: ClassId) {
+            if (rootNode.classId == classId) {
+                rootNode.inheritors?.clear()
+                return
+            }
+
+            rootNode.removeInheritor(classId)
+        }
+
+        fun isExhaustive() = containsNull && rootNode.isEmpty()
+    }
+
+    private data class SealedClassInheritors(val classId: ClassId, val inheritors: MutableSet<SealedClassInheritors>? = null) {
+        fun removeInheritor(classId: ClassId): Boolean {
+            return inheritors != null && (inheritors.removeIf { it.classId == classId } || inheritors.any { it.removeInheritor(classId) })
+        }
+
+        fun isEmpty(): Boolean {
+            return inheritors != null && inheritors.all { it.isEmpty() }
+        }
+    }
 
     private object SealedExhaustivenessVisitor : FirDefaultVisitor<Unit, SealedExhaustivenessData>() {
         override fun visitElement(element: FirElement, data: SealedExhaustivenessData) {}
@@ -171,7 +192,7 @@ class FirWhenExhaustivenessTransformer(private val bodyResolveComponents: BodyRe
         override fun visitResolvedTypeRef(resolvedTypeRef: FirResolvedTypeRef, data: SealedExhaustivenessData) {
             val lookupTag = (resolvedTypeRef.type as? ConeLookupTagBasedType)?.lookupTag ?: return
             val symbol = data.symbolProvider.getSymbolByLookupTag(lookupTag) as? FirClassSymbol ?: return
-            data.remainingInheritors.remove(symbol.classId)
+            data.removeInheritor(symbol.classId)
         }
 
         override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: SealedExhaustivenessData) {