FIR: Get rid of RealVariable.originalType

Store original type only for local variable aliases
The problem with the former is that RealVariable may be referenced
via safe-calls and via usual calls, but its type is fixed

^KT-38444 In Progress

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/FirDataFlowAnalyzer.kt

@@ -29,6 +29,7 @@ import org.jetbrains.kotlin.fir.types.*
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.name.Name
+import org.jetbrains.kotlin.utils.addIfNotNull
 
 class DataFlowAnalyzerContext<FLOW : Flow>(
     val graphBuilder: ControlFlowGraphBuilder,
@@ -118,8 +119,8 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         var variable = variableStorage.getRealVariableWithoutUnwrappingAlias(symbol, qualifiedAccessExpression, flow) ?: return null
         val result = mutableListOf<ConeKotlinType>()
         flow.directAliasMap[variable]?.let {
-            result += it.originalType
-            variable = it
+            result.addIfNotNull(it.originalType)
+            variable = it.variable
         }
         flow.getTypeStatement(variable)?.exactType?.let { result += it }
         return result.takeIf { it.isNotEmpty() }
@@ -783,7 +784,10 @@ abstract class FirDataFlowAnalyzer<FLOW : Flow>(
         }
 
         variableStorage.getOrCreateRealVariable(flow, initializer.symbol, initializer)?.let { initializerVariable ->
-            logicSystem.addLocalVariableAlias(flow, propertyVariable, initializerVariable)
+            logicSystem.addLocalVariableAlias(
+                flow, propertyVariable,
+                RealVariableAndType(initializerVariable, initializer.coneType)
+            )
             // node.flow.addImplication((propertyVariable notEq null) implies (initializerVariable notEq null))
         }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/LogicSystem.kt

@@ -15,12 +15,12 @@ abstract class Flow {
     abstract fun getVariablesInTypeStatements(): Collection<RealVariable>
     abstract fun removeOperations(variable: DataFlowVariable): Collection<Implication>
 
-    abstract val directAliasMap: Map<RealVariable, RealVariable>
+    abstract val directAliasMap: Map<RealVariable, RealVariableAndType>
     abstract val backwardsAliasMap: Map<RealVariable, List<RealVariable>>
 }
 
 fun Flow.unwrapVariable(variable: RealVariable): RealVariable {
-    return directAliasMap[variable] ?: variable
+    return directAliasMap[variable]?.variable ?: variable
 }
 
 abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceContext) {
@@ -53,7 +53,7 @@ abstract class LogicSystem<FLOW : Flow>(protected val context: ConeInferenceCont
         shouldRemoveSynthetics: Boolean,
     ): FLOW
 
-    abstract fun addLocalVariableAlias(flow: FLOW, alias: RealVariable, underlyingVariable: RealVariable)
+    abstract fun addLocalVariableAlias(flow: FLOW, alias: RealVariable, underlyingVariable: RealVariableAndType)
     abstract fun removeLocalVariableAlias(flow: FLOW, alias: RealVariable)
 
     protected abstract fun getImplicationsWithVariable(flow: FLOW, variable: DataFlowVariable): Collection<Implication>
@@ -190,4 +190,4 @@ fun <F : Flow> LogicSystem<F>.replaceVariableFromConditionInStatements(
         filter,
         transform,
     )
-}
+}

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/PersistentLogicSystem.kt

@@ -9,6 +9,7 @@ import com.google.common.collect.ArrayListMultimap
 import kotlinx.collections.immutable.*
 import org.jetbrains.kotlin.fir.types.ConeInferenceContext
 import org.jetbrains.kotlin.fir.types.ConeKotlinType
+import org.jetbrains.kotlin.utils.addIfNotNull
 import java.util.*
 import kotlin.NoSuchElementException
 
@@ -51,7 +52,7 @@ class PersistentFlow : Flow {
      * directAliasMap: { x -> a, y -> a}
      * backwardsAliasMap: { a -> [x, y] }
      */
-    override var directAliasMap: PersistentMap<RealVariable, RealVariable>
+    override var directAliasMap: PersistentMap<RealVariable, RealVariableAndType>
     override var backwardsAliasMap: PersistentMap<RealVariable, PersistentList<RealVariable>>
 
     constructor(previousFlow: PersistentFlow) {
@@ -156,9 +157,9 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
 
     private fun computeAliasesThatDontChange(
         flows: Collection<PersistentFlow>
-    ): MutableMap<RealVariable, RealVariable> {
+    ): MutableMap<RealVariable, RealVariableAndType> {
         val flowsSize = flows.size
-        val aliasedVariablesThatDontChangeAlias = mutableMapOf<RealVariable, RealVariable>()
+        val aliasedVariablesThatDontChangeAlias = mutableMapOf<RealVariable, RealVariableAndType>()
 
         flows.flatMapTo(mutableSetOf()) { it.directAliasMap.keys }.forEach { aliasedVariable ->
             val originals = flows.map { it.directAliasMap[aliasedVariable] ?: return@forEach }
@@ -173,7 +174,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
     }
 
     private fun PersistentFlow.addVariableAliases(
-        aliasedVariablesThatDontChangeAlias: MutableMap<RealVariable, RealVariable>
+        aliasedVariablesThatDontChangeAlias: MutableMap<RealVariable, RealVariableAndType>
     ) {
         for ((alias, underlyingVariable) in aliasedVariablesThatDontChangeAlias) {
             addLocalVariableAlias(this, alias, underlyingVariable)
@@ -189,11 +190,11 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
         }
     }
 
-    override fun addLocalVariableAlias(flow: PersistentFlow, alias: RealVariable, underlyingVariable: RealVariable) {
+    override fun addLocalVariableAlias(flow: PersistentFlow, alias: RealVariable, underlyingVariable: RealVariableAndType) {
         removeLocalVariableAlias(flow, alias)
         flow.directAliasMap = flow.directAliasMap.put(alias, underlyingVariable)
         flow.backwardsAliasMap = flow.backwardsAliasMap.put(
-            underlyingVariable,
+            underlyingVariable.variable,
             { persistentListOf(alias) },
             { variables -> variables + alias }
         )
@@ -201,7 +202,7 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
 
     override fun removeLocalVariableAlias(flow: PersistentFlow, alias: RealVariable) {
         flow.updatedAliasDiff += alias
-        val original = flow.directAliasMap[alias] ?: return
+        val original = flow.directAliasMap[alias]?.variable ?: return
         flow.directAliasMap = flow.directAliasMap.remove(alias)
         val variables = flow.backwardsAliasMap.getValue(original)
         flow.backwardsAliasMap = flow.backwardsAliasMap.put(original, variables - alias)
@@ -220,8 +221,8 @@ abstract class PersistentLogicSystem(context: ConeInferenceContext) : LogicSyste
                 flow = flow.previousFlow!!
             }
         } else {
-            result.exactType += variableUnderAlias.originalType
-            flow.approvedTypeStatements[variableUnderAlias]?.let { result += it }
+            result.exactType.addIfNotNull(variableUnderAlias.originalType)
+            flow.approvedTypeStatements[variableUnderAlias.variable]?.let { result += it }
         }
         return result
     }

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/VariableStorage.kt

@@ -41,7 +41,7 @@ class VariableStorage(private val session: FirSession) {
 
     private fun getOrCreateRealVariable(flow: Flow, symbol: AbstractFirBasedSymbol<*>, fir: FirElement): RealVariable {
         val variable = getOrCreateRealVariableWithoutUnwrappingAlias(flow, symbol, fir)
-        return flow.directAliasMap[variable] ?: variable
+        return flow.directAliasMap[variable]?.variable ?: variable
     }
 
     private fun FirElement.unwrapElement(): FirElement = when (this) {
@@ -85,13 +85,7 @@ class VariableStorage(private val session: FirSession) {
         }
 
         val receiverVariable = receiver?.let { getOrCreateVariable(flow, it) }
-        val originalType: ConeKotlinType = when (originalFir) {
-            is FirExpression -> originalFir.typeRef.coneTypeUnsafe()
-            is FirProperty -> originalFir.returnTypeRef.coneTypeUnsafe()
-            is FirVariableAssignment -> identifier.symbol.fir.extractReturnType()
-            else -> throw IllegalStateException("Should not be here: $originalFir")
-        }
-        return RealVariable(identifier, isThisReference, receiverVariable, originalType, counter++)
+        return RealVariable(identifier, isThisReference, receiverVariable, counter++)
     }
 
     @JvmName("getOrCreateRealVariableOrNull")

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/model.kt

@@ -40,7 +40,6 @@ class RealVariable(
     val identifier: Identifier,
     val isThisReference: Boolean,
     val explicitReceiverVariable: DataFlowVariable?,
-    val originalType: ConeKotlinType,
     variableIndexForDebug: Int
 ) : DataFlowVariable(variableIndexForDebug) {
     override val isStable: Boolean by lazy {
@@ -74,6 +73,26 @@ class RealVariable(
     }
 }
 
+class RealVariableAndType(val variable: RealVariable, val originalType: ConeKotlinType?) {
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (javaClass != other?.javaClass) return false
+
+        other as RealVariableAndType
+
+        if (variable != other.variable) return false
+        if (originalType != other.originalType) return false
+
+        return true
+    }
+
+    override fun hashCode(): Int {
+        var result = variable.hashCode()
+        result = 31 * result + originalType.hashCode()
+        return result
+    }
+}
+
 class SyntheticVariable(val fir: FirElement, variableIndexForDebug: Int) : DataFlowVariable(variableIndexForDebug) {
     override val isStable: Boolean get() = true
 
@@ -220,4 +239,4 @@ infix fun RealVariable.typeNotEq(type: ConeKotlinType): TypeStatement =
         MutableTypeStatement(this, linkedSetOf(), LinkedHashSet<ConeKotlinType>().apply { this += type })
     } else {
         MutableTypeStatement(this)
-    }
+    }