[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
2020-07-06 17:20:02 +03:00
parent 40cd30f7f6
commit 989e4293a3
5 changed files with 220 additions and 19 deletions
@@ -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)
+}
@@ -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))
+    }
@@ -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");
@@ -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");
@@ -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.FirRegularClass
-import org.jetbrains.kotlin.fir.declarations.collectEnumEntries
-import org.jetbrains.kotlin.fir.declarations.modality
+import org.jetbrains.kotlin.fir.declarations.*
 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 (inheritors.isNullOrEmpty()) return true
-        val data = SealedExhaustivenessData(
-            sealedClass.session.firSymbolProvider,
-            inheritors.toMutableSet(),
-            !nullable
-        )
+        if (sealedClass.sealedInheritors.isNullOrEmpty()) return true
+
+        val data = SealedExhaustivenessData(sealedClass, !nullable)
        for (branch in whenExpression.branches) {
            branch.condition.accept(SealedExhaustivenessVisitor, data)
        }
-        return data.containsNull && data.remainingInheritors.isEmpty()
+        return data.isExhaustive()
    }

-    private class SealedExhaustivenessData(
-        val symbolProvider: FirSymbolProvider,
-        val remainingInheritors: MutableSet<ClassId>,
-        var containsNull: Boolean
-    )
+    private inner class SealedExhaustivenessData(regularClass: FirRegularClass, var containsNull: Boolean) {
+        val symbolProvider = bodyResolveComponents.symbolProvider
+        private val rootNode = SealedClassInheritors(regularClass.classId, regularClass.sealedInheritors.mapToSealedInheritors())
+
+        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) {