[FIR] Avoid reanalyzing smartcast types in every when branch

The only new information that can appear among the smartcast types
when moving from one `when` branch to another is exactly what we already
collect into `flags`.

Author: lunakoly

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLDiagnosticsFe10TestGenerated.java

@@ -23,6 +23,7 @@ import org.jetbrains.kotlin.fir.resolve.fullyExpandedType
 import org.jetbrains.kotlin.fir.resolve.providers.symbolProvider
 import org.jetbrains.kotlin.fir.resolve.toRegularClassSymbol
 import org.jetbrains.kotlin.fir.resolve.toSymbol
+import org.jetbrains.kotlin.fir.resolve.transformers.WhenOnSealedClassExhaustivenessChecker.ConditionChecker.processBranch
 import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirClassSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirRegularClassSymbol
@@ -332,6 +333,12 @@ private object WhenOnBooleanExhaustivenessChecker : WhenExhaustivenessChecker()
         var containsFalse = false
     }
 
+    private fun recordValue(value: Any?, data: Flags) = when (value) {
+        true -> data.containsTrue = true
+        false -> data.containsFalse = true
+        else -> {}
+    }
+
     override fun computeMissingCases(
         whenExpression: FirWhenExpression,
         subjectType: ConeKotlinType,
@@ -345,6 +352,10 @@ private object WhenOnBooleanExhaustivenessChecker : WhenExhaustivenessChecker()
         }
 
         val flags = Flags()
+        (whenExpression.subjectVariable?.initializer as? FirSmartCastExpression)
+            ?.lowerTypesFromSmartCast
+            ?.mapNotNull { (it as? DfaType.BooleanLiteral)?.value }
+            ?.forEach { recordValue(it, flags) }
         whenExpression.accept(ConditionChecker, flags)
         if (!flags.containsTrue) {
             destination.add(WhenMissingCase.BooleanIsMissing.TrueIsMissing)
@@ -356,21 +367,10 @@ private object WhenOnBooleanExhaustivenessChecker : WhenExhaustivenessChecker()
 
     private object ConditionChecker : AbstractConditionChecker<Flags>() {
         override fun visitEqualityOperatorCall(equalityOperatorCall: FirEqualityOperatorCall, data: Flags) {
-            fun recordValue(value: Any?) = when (value) {
-                true -> data.containsTrue = true
-                false -> data.containsFalse = true
-                else -> {}
-            }
-
             if (equalityOperatorCall.operation.let { it == FirOperation.EQ || it == FirOperation.IDENTITY }) {
-                (equalityOperatorCall.arguments.firstOrNull() as? FirSmartCastExpression)
-                    ?.lowerTypesFromSmartCast
-                    ?.mapNotNull { (it as? DfaType.BooleanLiteral)?.value }
-                    ?.forEach(::recordValue)
-
                 val argument = equalityOperatorCall.arguments[1]
                 if (argument is FirLiteralExpression) {
-                    recordValue(argument.value)
+                    recordValue(argument.value, data)
                 }
             }
         }
@@ -393,6 +393,15 @@ private object WhenOnEnumExhaustivenessChecker : WhenExhaustivenessChecker() {
 
         val enumClass = (subjectType.toSymbol(session) as FirRegularClassSymbol).fir
         val notCheckedEntries = enumClass.declarations.mapNotNullTo(mutableSetOf()) { it as? FirEnumEntry }
+
+        whenExpression.subjectVariable?.initializer?.let { initializer ->
+            val knownNonValues = (initializer as? FirSmartCastExpression)
+                ?.lowerTypesFromSmartCast
+                ?.mapNotNull { (it as? DfaType.Symbol)?.symbol?.fir }
+                .orEmpty()
+            notCheckedEntries.removeAll(knownNonValues)
+        }
+
         whenExpression.accept(ConditionChecker, notCheckedEntries)
         notCheckedEntries.mapTo(destination) { WhenMissingCase.EnumCheckIsMissing(it.symbol.callableId) }
     }
@@ -402,12 +411,6 @@ private object WhenOnEnumExhaustivenessChecker : WhenExhaustivenessChecker() {
             if (!equalityOperatorCall.operation.let { it == FirOperation.EQ || it == FirOperation.IDENTITY }) return
             val argument = equalityOperatorCall.arguments[1]
 
-            val knownNonValues = (equalityOperatorCall.arguments.firstOrNull() as? FirSmartCastExpression)
-                ?.lowerTypesFromSmartCast
-                ?.mapNotNull { (it as? DfaType.Symbol)?.symbol?.fir }
-                .orEmpty()
-            data.removeAll(knownNonValues)
-
             @OptIn(UnsafeExpressionUtility::class)
             val symbol = argument.toResolvedCallableReferenceUnsafe()?.resolvedSymbol as? FirVariableSymbol<*> ?: return
             val checkedEnumEntry = symbol.fir as? FirEnumEntry ?: return
@@ -429,7 +432,12 @@ private object WhenOnSealedClassExhaustivenessChecker : WhenExhaustivenessChecke
     ) {
         val allSubclasses = subjectType.toSymbol(session)?.collectAllSubclasses(session) ?: return
         val checkedSubclasses = mutableSetOf<FirBasedSymbol<*>>()
-        whenExpression.accept(ConditionChecker, Flags(allSubclasses, checkedSubclasses, session))
+        val flags = Flags(allSubclasses, checkedSubclasses, session)
+
+        whenExpression.subjectVariable?.initializer?.let { initializer ->
+            inferVariantsFromSubjectSmartCast(initializer, flags)
+        }
+        whenExpression.accept(ConditionChecker, flags)
         (allSubclasses - checkedSubclasses).mapNotNullTo(destination) {
             when (it) {
                 is FirClassSymbol<*> -> WhenMissingCase.IsTypeCheckIsMissing(
@@ -449,14 +457,26 @@ private object WhenOnSealedClassExhaustivenessChecker : WhenExhaustivenessChecke
         val session: FirSession
     )
 
+    private fun inferVariantsFromSubjectSmartCast(subject: FirExpression, data: Flags) {
+        if (subject !is FirSmartCastExpression) return
+
+        for (knownNonType in subject.lowerTypesFromSmartCast) {
+            val symbol = when (knownNonType) {
+                is DfaType.Cone -> knownNonType.type.toSymbol(data.session)
+                is DfaType.Symbol -> knownNonType.symbol
+                else -> null
+            }
+            processBranch(symbol ?: continue, isNegated = false, data)
+        }
+    }
+
     private object ConditionChecker : AbstractConditionChecker<Flags>() {
         override fun visitEqualityOperatorCall(equalityOperatorCall: FirEqualityOperatorCall, data: Flags) {
             val isNegated = when (equalityOperatorCall.operation) {
                 FirOperation.EQ, FirOperation.IDENTITY -> false
                 FirOperation.NOT_EQ, FirOperation.NOT_IDENTITY -> true
                 else -> return
             }
-            inferVariantsFromSubjectSmartCast(equalityOperatorCall, data)
             val symbol = when (val argument = equalityOperatorCall.arguments[1].unwrapSmartcastExpression()) {
                 is FirResolvedQualifier -> {
                     val firClass = (argument.symbol as? FirRegularClassSymbol)?.fir
@@ -480,25 +500,11 @@ private object WhenOnSealedClassExhaustivenessChecker : WhenExhaustivenessChecke
                 FirOperation.NOT_IS -> true
                 else -> return
             }
-            inferVariantsFromSubjectSmartCast(typeOperatorCall, data)
             val symbol = typeOperatorCall.conversionTypeRef.coneType.fullyExpandedType(data.session).toSymbol(data.session) ?: return
             processBranch(symbol, isNegated, data)
         }
 
-        private fun inferVariantsFromSubjectSmartCast(typeOperatorCall: FirCall, data: Flags) {
-            val subject = typeOperatorCall.arguments.firstOrNull() as? FirSmartCastExpression ?: return
-
-            for (knownNonType in subject.lowerTypesFromSmartCast) {
-                val symbol = when (knownNonType) {
-                    is DfaType.Cone -> knownNonType.type.toSymbol(data.session)
-                    is DfaType.Symbol -> knownNonType.symbol
-                    else -> null
-                }
-                processBranch(symbol ?: continue, isNegated = false, data)
-            }
-        }
-
-        private fun processBranch(symbolToCheck: FirBasedSymbol<*>, isNegated: Boolean, flags: Flags) {
+        fun processBranch(symbolToCheck: FirBasedSymbol<*>, isNegated: Boolean, flags: Flags) {
             val subclassesOfType = symbolToCheck.collectAllSubclasses(flags.session)
             if (subclassesOfType.none { it in flags.allSubclasses }) {
                 return

analysis/low-level-api-fir/tests/org/jetbrains/kotlin/analysis/low/level/api/fir/diagnostic/compiler/based/LLReversedDiagnosticsFe10TestGenerated.java

@@ -0,0 +1,50 @@
+// FIR_IDENTICAL
+// RUN_PIPELINE_TILL: FRONTEND
+
+fun upperType(): Int {
+    var a: Boolean? = false
+    val block: () -> Unit = { a = null }
+
+    if (a == null) return 1
+    return <!NO_ELSE_IN_WHEN!>when<!> (a) {
+        true -> 2
+        false -> 3
+    }
+}
+
+fun booleanLiteral(): Int {
+    var a: Boolean = true
+    val block: () -> Unit = { a = false }
+
+    if (a == false) return 1
+    return <!NO_ELSE_IN_WHEN!>when<!> (a) {
+        true -> 2
+    }
+}
+
+enum class EnumBoolean { False, True }
+
+fun enumEntry(): Int {
+    var a: EnumBoolean = EnumBoolean.True
+    val block: () -> Unit = { a = EnumBoolean.False }
+
+    if (a == EnumBoolean.False) return 1
+    return <!NO_ELSE_IN_WHEN!>when<!> (a) {
+        EnumBoolean.True -> 2
+    }
+}
+
+sealed class SealedBoolean {
+    data object True : SealedBoolean()
+    data object False : SealedBoolean()
+}
+
+fun sealedVariant(): Int {
+    var a: SealedBoolean = SealedBoolean.True
+    val block: () -> Unit = { a = SealedBoolean.False }
+
+    if (a == SealedBoolean.False) return 1
+    return <!NO_ELSE_IN_WHEN!>when<!> (a) {
+        is SealedBoolean.True -> 2
+    }
+}