From 2e9cccd809a26fb66c1ca81fc05cf4b56224e20d Mon Sep 17 00:00:00 2001
From: pyos <pyos@google.com>
Date: Fri, 11 Nov 2022 18:57:11 +0100
Subject: [PATCH] FIR DFA: add a utility function for intersecting with
 original type

---
 .../fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt |  9 ++-------
 .../kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt |  7 +------
 .../org/jetbrains/kotlin/fir/resolve/dfa/util.kt  | 15 +++++++--------
 3 files changed, 10 insertions(+), 21 deletions(-)

diff --git a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt
index b23835fb0d9..e5541e11229 100644
--- a/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt
+++ b/compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/cfa/FirReturnsImpliesAnalyzer.kt
@@ -115,13 +115,8 @@ object FirReturnsImpliesAnalyzer : FirControlFlowChecker() {
 
         for ((realVar, requiredTypeStatement) in conditionStatements) {
             val fixedRealVar = typeStatements.keys.find { it.identifier == realVar.identifier } ?: realVar
-            val resultTypeStatement = typeStatements[fixedRealVar]
-
-            val resultType = mutableListOf<ConeKotlinType>().apply {
-                addIfNotNull(function.getParameterType(fixedRealVar.identifier.symbol, context))
-                if (resultTypeStatement != null) addAll(resultTypeStatement.exactType)
-            }.let { typeContext.intersectTypesOrNull(it) }
-
+            val originalType = function.getParameterType(fixedRealVar.identifier.symbol, context) ?: continue
+            val resultType = typeStatements[fixedRealVar]?.exactType.intersectWith(typeContext, originalType)
             val requiredType = typeContext.intersectTypesOrNull(requiredTypeStatement.exactType.toList())
             if (requiredType != null && !requiredType.isSupertypeOf(typeContext, resultType)) return true
         }
diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
index 8583d472c27..25d99af1aee 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt
@@ -102,12 +102,7 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
                 override fun receiverUpdated(symbol: FirBasedSymbol<*>, types: Set<ConeKotlinType>?) {
                     val index = receiverStack.getReceiverIndex(symbol) ?: return
                     val originalType = receiverStack.getOriginalType(index)
-                    val type = if (types?.isNotEmpty() == true) {
-                        typeContext.intersectTypes(types.toMutableList().also { it += originalType })
-                    } else {
-                        originalType
-                    }
-                    receiverStack.replaceReceiverType(index, type)
+                    receiverStack.replaceReceiverType(index, types.intersectWith(typeContext, originalType))
                 }
 
                 override val logicSystem: PersistentLogicSystem =
diff --git a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt
index e8b3e743a2f..a959fb389c5 100644
--- a/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt
+++ b/compiler/fir/semantics/src/org/jetbrains/kotlin/fir/resolve/dfa/util.kt
@@ -19,6 +19,7 @@ import org.jetbrains.kotlin.fir.symbols.impl.FirFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirNamedFunctionSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirSyntheticPropertySymbol
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.fir.types.ConeTypeContext
 import org.jetbrains.kotlin.fir.types.coneType
 import org.jetbrains.kotlin.name.StandardClassIds
 import kotlin.contracts.ExperimentalContracts
@@ -57,14 +58,12 @@ internal inline fun <K, V> PersistentMap<K, V>.put(
     }
 }
 
-@DfaInternals
-fun FirOperation.invert(): FirOperation = when (this) {
-    FirOperation.EQ -> FirOperation.NOT_EQ
-    FirOperation.NOT_EQ -> FirOperation.EQ
-    FirOperation.IDENTITY -> FirOperation.NOT_IDENTITY
-    FirOperation.NOT_IDENTITY -> FirOperation.IDENTITY
-    else -> throw IllegalArgumentException("$this can not be inverted")
-}
+fun Set<ConeKotlinType>?.intersectWith(context: ConeTypeContext, originalType: ConeKotlinType): ConeKotlinType =
+    if (!isNullOrEmpty()) {
+        context.intersectTypes(toMutableList().also { it += originalType })
+    } else {
+        originalType
+    }
 
 @DfaInternals
 fun FirOperation.isEq(): Boolean {